package io.modelcontextprotocol.autoconfigure.utils;


import cn.pinming.pmsuite.utils.StringUtils;
import io.modelcontextprotocol.autoconfigure.ToolCallback;
import io.modelcontextprotocol.autoconfigure.definition.ToolContext;
import io.modelcontextprotocol.server.McpServerFeatures;
import io.modelcontextprotocol.server.McpSyncServerExchange;
import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.McpSchema.Role;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;

import io.modelcontextprotocol.util.ListUtils;
import io.modelcontextprotocol.util.MapUtils;
import org.springframework.util.MimeType;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

public final class McpToolUtils {
    public static final String TOOL_CONTEXT_MCP_EXCHANGE_KEY = "exchange";

    private McpToolUtils() {
    }

    public static String prefixedToolName(String prefix, String toolName) {
        if (!StringUtils.isEmpty(prefix) && !StringUtils.isEmpty(toolName)) {
            String input = prefix + "_" + toolName;
            String formatted = input.replaceAll("[^a-zA-Z0-9_-]", "");
            formatted = formatted.replaceAll("-", "_");
            if (formatted.length() > 64) {
                formatted = formatted.substring(formatted.length() - 64);
            }

            return formatted;
        } else {
            throw new IllegalArgumentException("Prefix or toolName cannot be null or empty");
        }
    }

    public static List<McpServerFeatures.SyncToolSpecification> toSyncToolSpecification(List<ToolCallback> toolCallbacks) {
        return toolCallbacks.stream().map(McpToolUtils::toSyncToolSpecification).collect(Collectors.toList());
    }

    public static List<McpServerFeatures.SyncToolSpecification> toSyncToolSpecifications(ToolCallback... toolCallbacks) {
        return toSyncToolSpecification(ListUtils.of(toolCallbacks));
    }

    public static McpServerFeatures.SyncToolSpecification toSyncToolSpecification(ToolCallback toolCallback) {
        return toSyncToolSpecification(toolCallback, (MimeType)null);
    }

    public static McpServerFeatures.SyncToolSpecification toSyncToolSpecification(ToolCallback toolCallback, MimeType mimeType) {
        McpSchema.Tool tool = new McpSchema.Tool(toolCallback.getToolDefinition().name(), toolCallback.getToolDefinition().description(), toolCallback.getToolDefinition().inputSchema());
        return new McpServerFeatures.SyncToolSpecification(tool, (exchange, request) -> {
            try {
                String callResult = toolCallback.call(ModelOptionsUtils.toJsonString(request), new ToolContext(MapUtils.of("exchange", exchange)));
                return mimeType != null && mimeType.toString().startsWith("image") ? new McpSchema.CallToolResult(ListUtils.of(new McpSchema.ImageContent(ListUtils.of(Role.ASSISTANT), (Double)null, callResult, mimeType.toString())), false) : new McpSchema.CallToolResult(ListUtils.of(new McpSchema.TextContent(callResult)), false);
            } catch (Exception var5) {
                Exception e = var5;
                return new McpSchema.CallToolResult(ListUtils.of(new McpSchema.TextContent(e.getMessage())), true);
            }
        });
    }

    public static Optional<McpSyncServerExchange> getMcpExchange(ToolContext toolContext) {
        return toolContext != null && toolContext.getContext().containsKey("exchange") ? Optional.ofNullable((McpSyncServerExchange)toolContext.getContext().get("exchange")) : Optional.empty();
    }

    public static List<McpServerFeatures.AsyncToolSpecification> toAsyncToolSpecifications(List<ToolCallback> toolCallbacks) {
        return toolCallbacks.stream().map(McpToolUtils::toAsyncToolSpecification).collect(Collectors.toList());
    }

    public static List<McpServerFeatures.AsyncToolSpecification> toAsyncToolSpecifications(ToolCallback... toolCallbacks) {
        return toAsyncToolSpecifications(ListUtils.of(toolCallbacks));
    }

    public static McpServerFeatures.AsyncToolSpecification toAsyncToolSpecification(ToolCallback toolCallback) {
        return toAsyncToolSpecification(toolCallback, (MimeType)null);
    }

    public static McpServerFeatures.AsyncToolSpecification toAsyncToolSpecification(ToolCallback toolCallback, MimeType mimeType) {
        McpServerFeatures.SyncToolSpecification syncToolSpecification = toSyncToolSpecification(toolCallback, mimeType);
        return new McpServerFeatures.AsyncToolSpecification(syncToolSpecification.tool(), (exchange, map) -> {
            return Mono.fromCallable(() -> {
                return (McpSchema.CallToolResult)syncToolSpecification.call().apply(new McpSyncServerExchange(exchange), map);
            }).subscribeOn(Schedulers.boundedElastic());
        });
    }

}
